#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
External Tool Manager
外部工具管理器
"""

import os
import json
import time
import hashlib
import subprocess
import urllib.request
from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass
from pathlib import Path

@dataclass
class ToolInfo:
    """工具信息"""
    name: str
    version: str
    description: str
    url: str
    local_path: str
    checksum: str
    executable: bool = True
    dependencies: List[str] = None
    
    def __post_init__(self):
        if self.dependencies is None:
            self.dependencies = []

class ExternalToolManager:
    """外部工具管理器"""
    
    def __init__(self, tools_dir: str = "tools"):
        self.tools_dir = Path(tools_dir)
        self.tools_dir.mkdir(exist_ok=True)
        
        # 工具注册表
        self.tools_registry = {}
        
        # 初始化默认工具
        self._init_default_tools()
        
        # 工具状态缓存
        self.tool_status_cache = {}
        
        # 回调函数
        self.download_callbacks = {}
        self.execution_callbacks = {}
    
    def _init_default_tools(self):
        """初始化默认工具"""
        default_tools = [
            ToolInfo(
                name="linpeas",
                version="latest",
                description="Linux Privilege Escalation Awesome Script",
                url="https://github.com/carlospolop/PEASS-ng/releases/latest/download/linpeas.sh",
                local_path="linpeas.sh",
                checksum="",
                executable=True,
                dependencies=["bash", "curl"]
            ),
            ToolInfo(
                name="linenum",
                version="latest", 
                description="Linux Enumeration Script",
                url="https://raw.githubusercontent.com/rebootuser/LinEnum/master/LinEnum.sh",
                local_path="linenum.sh",
                checksum="",
                executable=True,
                dependencies=["bash"]
            ),
            ToolInfo(
                name="pspy",
                version="v1.2.1",
                description="Monitor linux processes without root permissions",
                url="https://github.com/DominicBreuker/pspy/releases/download/v1.2.1/pspy64",
                local_path="pspy64",
                checksum="",
                executable=True,
                dependencies=[]
            ),
            ToolInfo(
                name="gtfobins_db",
                version="latest",
                description="GTFOBins Database JSON",
                url="https://gtfobins.github.io/gtfobins.json",
                local_path="gtfobins.json",
                checksum="",
                executable=False,
                dependencies=[]
            )
        ]
        
        for tool in default_tools:
            self.register_tool(tool)
    
    def register_tool(self, tool_info: ToolInfo):
        """注册工具"""
        self.tools_registry[tool_info.name] = tool_info
        print(f"✅ 工具已注册: {tool_info.name} v{tool_info.version}")
    
    def unregister_tool(self, tool_name: str):
        """注销工具"""
        if tool_name in self.tools_registry:
            del self.tools_registry[tool_name]
            print(f"❌ 工具已注销: {tool_name}")
    
    def list_tools(self) -> List[Dict[str, Any]]:
        """列出所有工具"""
        tools_list = []
        
        for name, tool in self.tools_registry.items():
            status = self.get_tool_status(name)
            
            tools_list.append({
                "name": name,
                "version": tool.version,
                "description": tool.description,
                "status": status,
                "local_path": str(self.tools_dir / tool.local_path),
                "executable": tool.executable,
                "dependencies": tool.dependencies
            })
        
        return tools_list
    
    def get_tool_status(self, tool_name: str) -> str:
        """获取工具状态"""
        if tool_name not in self.tools_registry:
            return "not_registered"
        
        tool = self.tools_registry[tool_name]
        local_file = self.tools_dir / tool.local_path
        
        if not local_file.exists():
            return "not_downloaded"
        
        # 检查文件完整性
        if tool.checksum and not self._verify_checksum(local_file, tool.checksum):
            return "corrupted"
        
        # 检查可执行权限
        if tool.executable and not os.access(local_file, os.X_OK):
            return "not_executable"
        
        # 检查依赖
        if not self._check_dependencies(tool.dependencies):
            return "missing_dependencies"
        
        return "ready"
    
    def _verify_checksum(self, file_path: Path, expected_checksum: str) -> bool:
        """验证文件校验和"""
        try:
            with open(file_path, 'rb') as f:
                file_hash = hashlib.sha256(f.read()).hexdigest()
            return file_hash == expected_checksum
        except Exception:
            return False
    
    def _check_dependencies(self, dependencies: List[str]) -> bool:
        """检查依赖"""
        for dep in dependencies:
            try:
                subprocess.run(['which', dep], check=True, capture_output=True)
            except subprocess.CalledProcessError:
                return False
        return True
    
    def download_tool(self, tool_name: str, force: bool = False) -> bool:
        """下载工具"""
        if tool_name not in self.tools_registry:
            print(f"❌ 工具未注册: {tool_name}")
            return False
        
        tool = self.tools_registry[tool_name]
        local_file = self.tools_dir / tool.local_path
        
        # 检查是否已存在
        if local_file.exists() and not force:
            print(f"✅ 工具已存在: {tool_name}")
            return True
        
        try:
            print(f"📥 正在下载 {tool_name}...")
            
            # 下载文件
            urllib.request.urlretrieve(tool.url, local_file)
            
            # 设置可执行权限
            if tool.executable:
                os.chmod(local_file, 0o755)
            
            # 计算校验和
            if not tool.checksum:
                with open(local_file, 'rb') as f:
                    tool.checksum = hashlib.sha256(f.read()).hexdigest()
            
            print(f"✅ {tool_name} 下载完成")
            
            # 调用回调函数
            if tool_name in self.download_callbacks:
                self.download_callbacks[tool_name](tool_name, True, None)
            
            return True
            
        except Exception as e:
            print(f"❌ {tool_name} 下载失败: {e}")
            
            # 清理失败的下载
            if local_file.exists():
                local_file.unlink()
            
            # 调用回调函数
            if tool_name in self.download_callbacks:
                self.download_callbacks[tool_name](tool_name, False, str(e))
            
            return False
    
    def download_all_tools(self, force: bool = False) -> Dict[str, bool]:
        """下载所有工具"""
        results = {}
        
        for tool_name in self.tools_registry:
            results[tool_name] = self.download_tool(tool_name, force)
        
        return results
    
    def execute_tool(self, tool_name: str, args: List[str] = None, 
                    timeout: int = 300, capture_output: bool = True) -> Dict[str, Any]:
        """执行工具"""
        if tool_name not in self.tools_registry:
            return {
                "success": False,
                "error": f"工具未注册: {tool_name}",
                "stdout": "",
                "stderr": "",
                "returncode": -1
            }
        
        tool = self.tools_registry[tool_name]
        
        if not tool.executable:
            return {
                "success": False,
                "error": f"工具不可执行: {tool_name}",
                "stdout": "",
                "stderr": "",
                "returncode": -1
            }
        
        # 检查工具状态
        status = self.get_tool_status(tool_name)
        if status != "ready":
            return {
                "success": False,
                "error": f"工具状态异常: {status}",
                "stdout": "",
                "stderr": "",
                "returncode": -1
            }
        
        local_file = self.tools_dir / tool.local_path
        
        # 构造命令
        cmd = [str(local_file)]
        if args:
            cmd.extend(args)
        
        try:
            print(f"🚀 执行工具: {tool_name}")
            start_time = time.time()
            
            # 执行命令
            result = subprocess.run(
                cmd,
                timeout=timeout,
                capture_output=capture_output,
                text=True,
                cwd=self.tools_dir
            )
            
            execution_time = time.time() - start_time
            
            execution_result = {
                "success": result.returncode == 0,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "returncode": result.returncode,
                "execution_time": execution_time,
                "command": " ".join(cmd)
            }
            
            print(f"✅ {tool_name} 执行完成 (耗时: {execution_time:.2f}s)")
            
            # 调用回调函数
            if tool_name in self.execution_callbacks:
                self.execution_callbacks[tool_name](tool_name, execution_result)
            
            return execution_result
            
        except subprocess.TimeoutExpired:
            error_result = {
                "success": False,
                "error": f"执行超时 ({timeout}s)",
                "stdout": "",
                "stderr": "",
                "returncode": -1,
                "execution_time": timeout,
                "command": " ".join(cmd)
            }
            
            print(f"⏰ {tool_name} 执行超时")
            
            # 调用回调函数
            if tool_name in self.execution_callbacks:
                self.execution_callbacks[tool_name](tool_name, error_result)
            
            return error_result
            
        except Exception as e:
            error_result = {
                "success": False,
                "error": str(e),
                "stdout": "",
                "stderr": "",
                "returncode": -1,
                "execution_time": 0,
                "command": " ".join(cmd)
            }
            
            print(f"❌ {tool_name} 执行失败: {e}")
            
            # 调用回调函数
            if tool_name in self.execution_callbacks:
                self.execution_callbacks[tool_name](tool_name, error_result)
            
            return error_result
    
    def get_tool_path(self, tool_name: str) -> Optional[str]:
        """获取工具路径"""
        if tool_name not in self.tools_registry:
            return None
        
        tool = self.tools_registry[tool_name]
        local_file = self.tools_dir / tool.local_path
        
        if local_file.exists():
            return str(local_file)
        
        return None
    
    def remove_tool(self, tool_name: str) -> bool:
        """删除工具"""
        if tool_name not in self.tools_registry:
            return False
        
        tool = self.tools_registry[tool_name]
        local_file = self.tools_dir / tool.local_path
        
        try:
            if local_file.exists():
                local_file.unlink()
            
            print(f"🗑️ 工具已删除: {tool_name}")
            return True
            
        except Exception as e:
            print(f"❌ 删除工具失败: {e}")
            return False
    
    def update_tool(self, tool_name: str) -> bool:
        """更新工具"""
        return self.download_tool(tool_name, force=True)
    
    def update_all_tools(self) -> Dict[str, bool]:
        """更新所有工具"""
        return self.download_all_tools(force=True)
    
    def set_download_callback(self, tool_name: str, callback: Callable):
        """设置下载回调函数"""
        self.download_callbacks[tool_name] = callback
    
    def set_execution_callback(self, tool_name: str, callback: Callable):
        """设置执行回调函数"""
        self.execution_callbacks[tool_name] = callback
    
    def export_tools_info(self, file_path: str):
        """导出工具信息"""
        tools_info = {}
        
        for name, tool in self.tools_registry.items():
            tools_info[name] = {
                "version": tool.version,
                "description": tool.description,
                "url": tool.url,
                "local_path": tool.local_path,
                "checksum": tool.checksum,
                "executable": tool.executable,
                "dependencies": tool.dependencies,
                "status": self.get_tool_status(name)
            }
        
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(tools_info, f, indent=2, ensure_ascii=False)
        
        print(f"📄 工具信息已导出到: {file_path}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取统计信息"""
        total_tools = len(self.tools_registry)
        ready_tools = 0
        downloaded_tools = 0
        
        status_counts = {}
        
        for tool_name in self.tools_registry:
            status = self.get_tool_status(tool_name)
            
            if status == "ready":
                ready_tools += 1
                downloaded_tools += 1
            elif status in ["not_executable", "missing_dependencies", "corrupted"]:
                downloaded_tools += 1
            
            status_counts[status] = status_counts.get(status, 0) + 1
        
        return {
            "total_tools": total_tools,
            "ready_tools": ready_tools,
            "downloaded_tools": downloaded_tools,
            "ready_percentage": (ready_tools / total_tools * 100) if total_tools > 0 else 0,
            "status_breakdown": status_counts
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 创建工具管理器
    tool_manager = ExternalToolManager("./external_tools")
    
    # 列出所有工具
    print("📋 已注册的工具:")
    tools = tool_manager.list_tools()
    for tool in tools:
        print(f"  - {tool['name']}: {tool['status']}")
    
    # 下载所有工具
    print("\n📥 下载工具...")
    download_results = tool_manager.download_all_tools()
    
    for tool_name, success in download_results.items():
        if success:
            print(f"✅ {tool_name} 下载成功")
        else:
            print(f"❌ {tool_name} 下载失败")
    
    # 执行LinPEAS
    print("\n🚀 执行LinPEAS...")
    result = tool_manager.execute_tool("linpeas", ["-a"])
    
    if result["success"]:
        print("✅ LinPEAS执行成功")
        print(f"输出长度: {len(result['stdout'])} 字符")
    else:
        print(f"❌ LinPEAS执行失败: {result['error']}")
    
    # 获取统计信息
    stats = tool_manager.get_statistics()
    print(f"\n📊 统计信息:")
    print(f"  总工具数: {stats['total_tools']}")
    print(f"  就绪工具数: {stats['ready_tools']}")
    print(f"  就绪率: {stats['ready_percentage']:.1f}%")

if __name__ == "__main__":
    example_usage()